/* Copyright (c) 2021, National University of Defense Technology. All rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 * 
 *     http://www.apache.org/licenses/LICENSE-2.0
 * 
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include <cmath>
#include <iostream>
#include "vec_dot.h"
#include "parvec_axpby.h"
#include "par_pipecg_solver.h"

namespace YHAMG
{

int PCG_RR_STEPS = 50;

Par_PipeCG_Solver::Par_PipeCG_Solver(const Par_Operator& _A)
	: A(&_A),
	P(0),
	MaxIters(500),
	Tolerance(1.0e-08),
	PrintStats(0)
{
}

Par_PipeCG_Solver::Par_PipeCG_Solver(const Par_Operator& _A, const Par_Operator& _P)
	: A(&_A),
	P(&_P),
	MaxIters(500),
	Tolerance(1.0e-08),
	PrintStats(0)
{
}

Par_PipeCG_Solver::Par_PipeCG_Solver(const Par_Operator& _A, int max_Iters, double tolerance, int print_stats)
	: A(&_A),
	P(0),
	MaxIters(max_Iters),
	Tolerance(tolerance),
	PrintStats(print_stats)
{
}

Par_PipeCG_Solver::Par_PipeCG_Solver(const Par_Operator& _A, const Par_Operator& _P, int max_Iters, double tolerance, int print_stats)
	: A(&_A),
	P(&_P),
	MaxIters(max_Iters),
	Tolerance(tolerance),
	PrintStats(print_stats)
{
}

#define APPLYDOT(x, y, result) ApplyDot(x, y, result, buf, dst, cnt)
#define BEGINREDUCTION()       BeginReduction(buf, cnt, comm, &request)
#define ENDREDUCTION()         EndReduction(buf, dst, cnt, &request, &status)

static inline void ApplyDot(const Par_Vector& x, const Par_Vector& y, double& result, double* buf, double **dst, int &cnt)
{
	buf[cnt] = Vec_Dot(x.Loc, y.Loc);
	dst[cnt++] = &result;
}

static inline void BeginReduction(double* buf, int cnt, MPI_Comm comm, MPI_Request* request)
{
	MPI_Iallreduce(MPI_IN_PLACE, buf, cnt, MPI_DOUBLE, MPI_SUM, comm, request);
}

static inline void EndReduction(const double* buf, double** dst, int &cnt, MPI_Request* request, MPI_Status* status)
{
	MPI_Wait(request, status);
	for (int i = 0; i < cnt; ++i)
		*dst[i] = buf[i];
	cnt = 0;
}

void Par_PipeCG_Solver::operator()(const Par_Vector& b, const Par_Vector& x, int& iter, double& relres) const
{
	MPI_Comm comm = A->comm;

	int rank;
	MPI_Comm_rank(comm, &rank);

	MPI_Request request;
	MPI_Status status;
	double buf[3];
	double* dst[3];
	int cnt = 0;

	int n = A->OutSize();
	
	double normb, res, alpha, beta, omega, rho1, rho;
	Par_Vector r(comm), s(comm), t(comm), p(comm), q(comm), u(comm), v(comm), w(comm), z(comm);

	r.Resize(n);
	s.Resize(n);
	t.Resize(n);
	p.Resize(n);
	q.Resize(n);
	u.Resize(n);
	v.Resize(n);
	w.Resize(n);
	z.Resize(n);

	if (PrintStats && rank == 0)
	{
		std::cout << "Iter\tResidual\n";
		std::cout << "----------------------------------------------------------------\n";
	}

	iter = 0;

	APPLYDOT(b, b, normb);

	BEGINREDUCTION();

	A->Apply(x, r);
	ParVec_AXPBY(-1.0, r, 1.0, b, r);

	if (P)
		P->Apply(r, u);
	else
		u.Copy(r);

	A->Apply(u, w);

	ENDREDUCTION();

	normb = sqrt(normb);

	while (1)
	{
		APPLYDOT(r, u, rho);
		APPLYDOT(w, u, omega);
		APPLYDOT(r, r, res);

		BEGINREDUCTION();

		if (P)
			P->Apply(w, t);
		else
			t.Copy(w);

		A->Apply(t, v);

		ENDREDUCTION();

		res = sqrt(res);

		if (PrintStats && rank == 0)
			std::cout << iter << "\t" << res << "\n";

		if (res / normb <= Tolerance || iter == MaxIters) break;

		if (iter == 0)
		{
			alpha = rho / omega;

			z.Copy(v);
			q.Copy(t);
			s.Copy(w);
			p.Copy(u);
		}
		else
		{
			beta = rho / rho1;
			alpha = 1 / (omega / rho - beta / alpha);

			ParVec_AXPBY(1.0, v, beta, z, z);
			ParVec_AXPBY(1.0, t, beta, q, q);
			ParVec_AXPBY(1.0, w, beta, s, s);
			ParVec_AXPBY(1.0, u, beta, p, p);
		}

		x.AddScaled(alpha, p);
		r.AddScaled(-alpha, s);
		u.AddScaled(-alpha, q);
		w.AddScaled(-alpha, z);

		rho1 = rho;

		++iter;

		if (iter % PCG_RR_STEPS == 0)
		{
			A->Apply(x, r);
			ParVec_AXPBY(-1.0, r, 1.0, b, r);
			if (P)
				P->Apply(r, u);
			else
				u.Copy(r);
			A->Apply(u, w);
		}
	}

	relres = res / normb;

	if (PrintStats && rank == 0)
	{
		std::cout << "----------------------------------------------------------------\n";
		std::cout << "Iterations: " << iter << "\n";
		std::cout << "Final Relative Residual: " << relres << "\n";
	}
}

}